import os
import torch
import torch.distributed as dist
from torch.multiprocessing import spawn

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29503'
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train(rank, world_size):
    setup(rank, world_size)
    print(f"Rank {rank}: Before barrier")
    dist.barrier()
    print(f"Rank {rank}: After barrier")
    tensor = torch.ones(1).cuda(rank) * rank
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    print(f"Rank {rank}: Tensor value after all_reduce is {tensor.item()}")
    dist.destroy_process_group()

if __name__ == "__main__":
    world_size = torch.cuda.device_count()
    spawn(train, args=(world_size,), nprocs=world_size, join=True)
